package dsg

import (
	"fmt"
	"strings"
)

type SAvlNode struct {
	value     Value
	left      int
	right     int
	bf        int
}

type SAvlTree struct {
	root      int
	comp_func CompFunc

      node     []SAvlNode
      ls_node  *LinkSet
}

func InitSAvlTree(comp_func CompFunc, max_node int) *SAvlTree {
      var tree SAvlTree
      tree.root = -1
      tree.comp_func = comp_func
      tree.node = make ([]SAvlNode, max_node)
      tree.ls_node = InitFullLinkSet (max_node)
	return &tree
}

func (tr *SAvlTree) Add(v Value) {
	if tr.root == -1 {
		tr.root = tr.InitSAvlNode(v)
	} else {
		tr.root, _ = tr.AddNode (tr.root, v)
	}
}

func (tr *SAvlTree) Remove(v Value) {
	if tr.root != -1 {
            var junk int
            tr.root, junk = tr.RemoveNode(tr.root, v)
            if junk != -1 {
                  tr.ls_node.Set (junk)
            }
	}
}

func (tr *SAvlTree) Search (v Value, dir []int) int {
	if tr.root != -1 {
		return tr.SearchNode (tr.root, v, dir)
	} else {
            return 0
      }
}

func (tr *SAvlTree) FollowPath (path []int) int {
	l := len(path)
	node := tr.root
	if node == -1 {return -1}
	for i := 0; i < l; i ++ {
		switch path[i] {
		case -1:
			node = tr.node[node].left
			if node == -1 {return -1}
		case 1:
			node = tr.node[node].right
			if node == -1 {return -1}
		case 0:
			break
		}
	}
	return node
}

func (tr *SAvlTree) SearchRange (v Value, work_path []int) (vc Value, vl Value, vr Value) {
	if tr.root == -1 { return } 
      l := tr.SearchNode (tr.root, v, work_path)
      left_neigh, right_neigh := -1, -1

	if work_path[l - 1] == 0 {
		node := tr.FollowPath (work_path[:l - 1])
		vc = tr.node[node].value
		right_neigh = tr.FindRightNeighbour(node)
		left_neigh = tr.FindLeftNeighbour(node)
	} 

	if left_neigh == -1 {
		sl := -1
		for i := l - 1; i >= 0; i -- {
			if work_path[i] == 1 {
				sl = i
				break
			}
		}
		if (sl != -1) {
			left_neigh = tr.FollowPath (work_path[:sl])
		}
	}

	if right_neigh == -1 {
		sr := -1
		for i := l - 1; i >= 0; i -- {
			if work_path[i] == -1 {
				sr = i
				break
			}
		}
		if (sr != -1) {
			right_neigh = tr.FollowPath (work_path[:sr])
		}
	}

	if left_neigh != -1 {
		vl = tr.node[left_neigh].value
	} 
	if right_neigh != -1 {
		vr = tr.node[right_neigh].value
	} 

	return 
}

func (tree * SAvlTree) InitSAvlNode (v Value) int {
      id_node := tree.ls_node.GetFirstLabel ()
      if id_node == -1 {
            panic ("not enough space for SAvlTree nodes")
      }
      tree.node[id_node] = SAvlNode {v, -1, -1, 0}
      tree.ls_node.UnSet (id_node)
	return id_node
}

func (tr * SAvlTree) SearchNode (root int, v Value, dir []int) int {
      rootn := &tr.node[root]
	comp_res := tr.comp_func (v, rootn.value)
      dir[0] = comp_res
      var ndir int
	if comp_res < 0 && rootn.left != -1 {
            ndir = tr.SearchNode (rootn.left, v, dir[1:])
	} else if comp_res > 0 && rootn.right != -1 {
            ndir = tr.SearchNode (rootn.right, v, dir[1:])
	}
      return ndir + 1
}

func (tr * SAvlTree) FindRightNeighbour (root int) int {
	if root == -1 { return -1}
	node := tr.node[root].right
	if node == -1 {return -1}
	for ; tr.node[node].left != -1; node = tr.node[node].left {}
	return node
}

func (tr * SAvlTree) FindLeftNeighbour (root int) int {
	if root == -1 { return -1}
	node := tr.node[root].left
	if node == -1 {return -1}
	for ; tr.node[node].right != -1; node = tr.node[node].right {}
	return node
}


func (tree *SAvlTree) AddNode(root int, v Value) (new_root, new_deeper int) {
	var deeper int
      rootn := &tree.node[root]

	if tree.comp_func(v, rootn.value) < 0 {
		if rootn.left == -1 {
			rootn.left = tree.InitSAvlNode(v)
		} else {
			rootn.left, deeper = tree.AddNode(rootn.left, v)
			if deeper == 0 {
				return root, 0
			}
		}
		if rootn.bf == -1 {
			new_root = tree.BalanceLeft(root)
			new_deeper = 0
			return
		}
		rootn.bf--
		new_deeper = rootn.bf
		new_root = root
		return
	}

	if rootn.right == -1 {
		rootn.right = tree.InitSAvlNode(v)
	} else {
		rootn.right, deeper = tree.AddNode (rootn.right, v)
		if deeper == 0 {
			return root, 0
		}
	}
	if rootn.bf == 1 {
		new_root = tree.BalanceRight(root)
		new_deeper = 0
		return
	}
	rootn.bf++
	new_deeper = rootn.bf
	new_root = root
	return
}

func (tr * SAvlTree) BalanceLeft (root int) int {
      var p, q int
	var pn, qn *SAvlNode
      rootn := &tr.node[root]

      if tr.node[rootn.left].bf == 0 {
            p = rootn.left
            pn = &tr.node[p]
            rootn.left = pn.right
            pn.right = root
            rootn.bf = -1
            pn.bf = 1
            return p
      }

      if tr.node[rootn.left].bf < 0 {
            p = rootn.left
            pn = &tr.node[p]
		rootn.left = pn.right
		rootn.bf = 0
	} else {
            q = rootn.left
		qn = &tr.node[q]
            p = qn.right
		pn = &tr.node[p]
		qn.right = pn.left
		rootn.left = pn.right
		pn.left = q
		rootn.bf = 0
		qn.bf = 0
		if pn.bf == -1 {
			rootn.bf = 1
		}
		if pn.bf == 1 {
			qn.bf = -1
		}
	}
	pn.right = root
	pn.bf = 0
	return p
}

func (tr * SAvlTree) BalanceRight (root int) int {
      var p, q int
	var pn, qn *SAvlNode
      rootn := &tr.node[root]

      if tr.node[rootn.right].bf == 0 {
            p = rootn.right
            pn = &tr.node[p]
            rootn.right = pn.left
            pn.left = root
            rootn.bf = 1
            pn.bf = -1
            return p
      }

      if tr.node[rootn.right].bf > 0 {
            p = rootn.right
            pn = &tr.node[p]
		rootn.right = pn.left
		rootn.bf = 0
	} else {
            q = rootn.right
		qn = &tr.node[q]
            p = qn.left
		pn = &tr.node[p]
		qn.left = pn.right
		rootn.right = pn.left
		pn.right = q
		rootn.bf = 0
		qn.bf = 0
		if pn.bf == 1 {
			rootn.bf = -1
		}
		if pn.bf == -1 {
			qn.bf = 1
		}
	}
	pn.left = root
	pn.bf = 0
	return p
}

func (tr * SAvlTree) RemoveNode (root int, v Value) (new_root, junk int) {
      rootn := &tr.node[root]
	comp_res := tr.comp_func(v, rootn.value)
	if comp_res == 0 {
		junk = root
		if rootn.right == -1 {
			return rootn.left, junk
		}
		oldbf := tr.node[rootn.right].bf
		rootn.right, new_root = tr.RemoveLeftMostDescendant(rootn.right)
            new_root_n := &tr.node[new_root]
		new_root_n.left = rootn.left
		new_root_n.right = rootn.right
		new_root_n.bf = rootn.bf
		new_root = tr.RestoreRightBalance(new_root, oldbf)
		return
	} else if comp_res < 0 {
		if rootn.left == -1 {
			return root, -1
		}
		oldbf := tr.node[rootn.left].bf
		rootn.left, junk = tr.RemoveNode(rootn.left, v)
		new_root = tr.RestoreLeftBalance(root, oldbf)
		return
	} else {
		if rootn.right == -1 {
			return root, -1
		}
		oldbf := tr.node[rootn.right].bf
		rootn.right, junk = tr.RemoveNode(rootn.right, v)
		new_root = tr.RestoreRightBalance(root, oldbf)
		return
	}
}

func (tr * SAvlTree) RemoveLeftMostDescendant(root int) (new_root, junk int) {
      rootn := &tr.node[root]
	left_child := rootn.left
	if left_child == -1 {
		junk = root
		new_root = rootn.right
		return
	}
	oldbf := tr.node[left_child].bf
	rootn.left, junk = tr.RemoveLeftMostDescendant(left_child)
	new_root = tr.RestoreLeftBalance(root, oldbf)
	return
}

func (tr * SAvlTree) RestoreLeftBalance (root int, oldbf int) int {
      rootn := &tr.node[root]
	left_child := rootn.left
	if left_child == -1 {
		rootn.bf++
	} else if tr.node[left_child].bf != oldbf && tr.node[left_child].bf == 0 {
		rootn.bf++
	}
	if rootn.bf > 1 {
		return tr.BalanceRight(root)
	}
	return root
}

func (tr * SAvlTree) RestoreRightBalance (root int, oldbf int) int {
      rootn := &tr.node[root]
	right_child := rootn.right
	if right_child == -1 {
		rootn.bf--
	} else if tr.node[right_child].bf != oldbf && tr.node[right_child].bf == 0 {
		rootn.bf--
	}
	if rootn.bf < -1 {
		return tr.BalanceLeft(root)
	}
	return root
}

func (tr * SAvlTree) GetHeight () int {
      if tr.root == -1 {
            return 0
      } else {
            return tr.GetHeightNode (tr.root)
      }
}

func (tr * SAvlTree) GetHeightNode (root int) int {
	var hl, hr int
      rootn := &tr.node[root]

	if rootn.left != -1 {
		hl = tr.GetHeightNode(rootn.left)
	}
	if rootn.right != -1 {
		hr = tr.GetHeightNode(rootn.right)
	}
	var height int
	if hl < hr {
		height = hr
	} else {
		height = hl
	}

	return height + 1
}

func (tr * SAvlTree) PrintNode (root, level int) {
	prefix := strings.Repeat(" ", level * 10)
	prefix_new := strings.Repeat(" ", (level + 1) * 10)
      rootn := &tr.node[root]

	fmt.Printf("%s%8d%+2d\n", prefix, rootn.value.(int), rootn.bf)
	if rootn.left == -1 {
		fmt.Printf("%s%s\n", prefix_new, "--NoLeft--")
	} else {
            tr.PrintNode (rootn.left, level + 1)
	}

	if rootn.right == -1 {
		fmt.Printf("%s%s\n", prefix_new, "--NoRight-")
	} else {
            tr.PrintNode (rootn.right, level + 1)
	}
}

func (tr *SAvlTree) Print() {
	if tr.root == -1 {
		fmt.Print("--Empty---\n")
	} else {
            tr.PrintNode (tr.root, 0)
	}
}
